{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "07be5a73-b9b6-4742-9cb7-4685ac6635a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import cv2\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import torchvision\n",
    "import numpy as np\n",
    "from IPython.display import  Image, display,clear_output\n",
    "device = torch.device(\"cuda\",0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee4cd287-afb5-4496-b7c4-ff2d4d7bff7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in /home/robotnx/.cache/torch/hub/ultralytics_yolov5_master\n",
      "YOLOv5 🚀 2021-10-15 torch 1.8.0 CUDA:0 (Xavier, 7765.4140625MB)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "model = torch.hub.load('ultralytics/yolov5', 'yolov5s')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b16cd50-96ab-4a57-a539-d9f8eff136f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "video = cv2.VideoCapture(\"/dev/video1\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "447a94d7-6277-4e4d-9c46-864e6419ad27",
   "metadata": {},
   "outputs": [],
   "source": [
    "ret,frame = video.read()\n",
    "print(ret)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f776b39b-b9a7-4d01-a716-4a9a7b74de2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import threading\n",
    "import time\n",
    "\n",
    "exitFlag = 0\n",
    "\n",
    "class VideoPlay (threading.Thread):\n",
    "    def __init__(self, threadID, name, device):\n",
    "        threading.Thread.__init__(self)\n",
    "        self.threadID = threadID\n",
    "        self.name = name\n",
    "        self.device = device\n",
    "        self.isRuning = True\n",
    "        self.display_handle=display(None, display_id=True)\n",
    "        self.video = device\n",
    "    def stop(self):\n",
    "        self.isRuning = False\n",
    "    def readOneFrame(self):\n",
    "        flag,self.currentFrame = self.video.read()\n",
    "        return self.currentFrame\n",
    "    def show(self):\n",
    "        self.display_handle.update(Image(data=cv2.imencode(\".jpeg\", self.currentFrame)[1].tobytes()))\n",
    "    def run(self):\n",
    "        if self.isOpened:\n",
    "            while(self.isRuning):\n",
    "                display_handle.update(Image(data=cv2.imencode(\".jpeg\", frame)[1].tobytes()))\n",
    "                flag, self.currentFrame =self.video.read()\n",
    "                time.sleep(0.02)\n",
    "        video.release()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37d9b4cb-a23f-4e6f-98ba-c119416b1d6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "videoPlay = VideoPlay(1,\"camePlay\",video)\n",
    "half = False\n",
    "half &= device.type != 'cpu'\n",
    "while(True):\n",
    "    frame = videoPlay.readOneFrame()\n",
    "    # img = torch.from_numpy(frame)\n",
    "    # img = img.half() if half else img.float()  # uint8 to fp16/32\n",
    "    # img = img.transfor\n",
    "    # img /= 255.0  # 0 - 255 to 0.0 - 1.0\n",
    "    # img = img.unsqueeze(0)\n",
    "    # pred = model(img)\n",
    "    pred = model(frame)\n",
    "    for r in pred.pred[0]:\n",
    "        ir = r.int()\n",
    "        # print(pred.names[ir[5]])\n",
    "        leftP =  (ir[0].item(),ir[1].item())\n",
    "        rightP = (ir[2].item(),ir[3].item())\n",
    "        centerx =  int((leftP[0] + rightP[0] )/ 2 - 10)\n",
    "        centery =  int((leftP[1] + rightP[1] )/ 2)\n",
    "        cv2.rectangle(frame,leftP,rightP,(0, 255, 0), 2)\n",
    "        string =  pred.names[ir[5]] \n",
    "        cv2.putText(frame,string,(centerx, centery), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (255, 255, 255), 2)\n",
    "    videoPlay.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99c88f69-18ea-49f8-acc0-56b169d6e901",
   "metadata": {},
   "outputs": [],
   "source": [
    "videoPlay.setUp()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
